import paho.mqtt.client as mqtt
import time
import json
import threading
from typing import Dict, List, Optional, Tuple, Any, Callable

import random

class MQTTTestLib:
    def __init__(self, broker_host: str = "localhost", broker_port: int = 1883, keepalive: int = 60,
                 msg_cache_max_size: int = 10000, msg_cache_expire: int = 3600):
        """
        初始化MQTT测试工具库
        :param broker_host: MQTT Broker的主机地址，默认值为"localhost"
        :param broker_port: MQTT Broker的端口号，默认值为1883
        :param keepalive: 客户端与Broker的保活时间（秒），默认值为60
        :param msg_cache_max_size: 消息缓存的最大条数限制，0表示无限制，默认值为10000
        :param msg_cache_expire: 消息缓存的过期时间（秒），0表示永不过期，默认值为3600
        """
        self.broker_host = broker_host
        self.broker_port = broker_port
        self.keepalive = keepalive

        # 客户端存储：key为客户端ID，value为mqtt.Client对象
        self.clients: Dict[str, mqtt.Client] = {}
        # 客户端连接状态：key为客户端ID，value为连接状态（True表示已连接）
        self.client_status: Dict[str, bool] = {}

        # 消息缓存配置
        self.msg_cache_max_size = msg_cache_max_size
        self.msg_cache_expire = msg_cache_expire
        # 消息缓存存储：key为主题，value为消息列表[(payload, timestamp), ...]
        self.msg_cache: Dict[str, List[Tuple[str, float]]] = {}
        # 缓存操作锁，确保线程安全
        self.msg_cache_lock = threading.Lock()

        # 异步发布相关
        self.async_pub_threads: List[threading.Thread] = []  # 异步发布线程列表
        self.async_pub_threads_lock = threading.Lock()  # 线程列表操作锁
        # 异步发布回调函数：(client_id, topic, success, error_msg) -> None
        self.async_pub_callback: Optional[Callable[[str, str, bool, str], None]] = None

    def _on_connect(self, client: mqtt.Client, userdata: Any, flags: Dict, rc: int, properties: Any = None):
        """
        内部连接回调函数（v5.0 API版本）

        :param client: 触发回调的mqtt.Client对象
        :param userdata: 客户端的用户数据
        :param flags: 响应标志字典
        :param rc: 连接结果代码，0表示成功
        :param properties: MQTT v5.0属性（可选）
        """
        # 从客户端ID获取连接状态键
        client_id = client._client_id.decode() if isinstance(client._client_id, bytes) else client._client_id
        self.client_status[client_id] = (rc == 0)

    def _on_message(self, client: mqtt.Client, userdata: Any, msg: mqtt.MQTTMessage):
        """
        内部消息接收回调函数（v5.0 API版本）

        :param client: 触发回调的mqtt.Client对象
        :param userdata: 客户端的用户数据
        :param msg: 接收到的消息对象，包含topic和payload属性
        """
        topic = msg.topic
        payload = msg.payload.decode()  # 将字节流解码为字符串
        timestamp = time.time()  # 记录接收时间戳

        with self.msg_cache_lock:  # 线程安全操作
            # 清理过期消息
            if self.msg_cache_expire > 0 and topic in self.msg_cache:
                self.msg_cache[topic] = [
                    (p, ts) for p, ts in self.msg_cache[topic]
                    if (timestamp - ts) <= self.msg_cache_expire
                ]

            # 存储新消息
            if topic not in self.msg_cache:
                self.msg_cache[topic] = []
            self.msg_cache[topic].append((payload, timestamp))

            # 处理缓存大小限制
            if self.msg_cache_max_size > 0 and len(self.msg_cache[topic]) > self.msg_cache_max_size:
                excess = len(self.msg_cache[topic]) - self.msg_cache_max_size
                self.msg_cache[topic] = self.msg_cache[topic][excess:]  # 保留最新的消息

    def create_client(self, client_id: str, clean_session: bool = True,
                      username: Optional[str] = None, password: Optional[str] = None,
                      will_topic: Optional[str] = None, will_payload: Optional[str] = None,
                      will_qos: int = 0) -> mqtt.Client:
        """
        创建MQTT客户端并配置相关参数

        :param client_id: 客户端唯一标识符，不可重复
        :param clean_session: 是否使用干净会话，True表示断开后清除会话状态，默认值为True
        :param username: 认证用户名（可选）
        :param password: 认证密码（可选）
        :param will_topic: 遗嘱消息主题（可选）
        :param will_payload: 遗嘱消息内容（可选，需与will_topic同时设置）
        :param will_qos: 遗嘱消息的QoS级别，默认值为0
        :return: 创建成功的mqtt.Client对象
        """
        # 如果客户端已存在，先断开连接
        if client_id in self.clients:
            self.disconnect_client(client_id)

        # 创建客户端并指定使用v5回调API，消除 deprecation 警告
        client = mqtt.Client(client_id=client_id, clean_session=clean_session, callback_api_version=mqtt.CallbackAPIVersion.VERSION2)
        
        # 绑定回调函数
        client.on_connect = self._on_connect
        client.on_message = self._on_message

        # 配置认证信息
        if username and password:
            client.username_pw_set(username, password)

        # 配置遗嘱消息
        if will_topic and will_payload:
            client.will_set(will_topic, will_payload, will_qos)

        # 存储客户端并初始化状态
        self.clients[client_id] = client
        self.client_status[client_id] = False
        return client

    def connect_client(self, client_id: str, timeout: int = 5) -> bool:
        """
        连接客户端到MQTT Broker

        :param client_id: 已创建的客户端ID
        :param timeout: 连接超时时间（秒），默认值为5
        :return: 连接成功返回True，超时或失败返回False
        :raises ValueError: 客户端ID不存在时抛出
        :raises ConnectionError: 连接过程中发生异常时抛出
        """
        if client_id not in self.clients:
            raise ValueError(f"客户端 '{client_id}' 未创建，请先调用 create_client 方法")

        client = self.clients[client_id]
        start_time = time.time()

        try:
            # 发起连接
            client.connect(self.broker_host, self.broker_port, self.keepalive)
            # 启动后台消息处理循环
            client.loop_start()

            # 等待连接结果
            while time.time() - start_time < timeout:
                if self.client_status[client_id]:
                    return True
                time.sleep(0.1)  # 降低CPU占用
            return False

        except Exception as e:
            raise ConnectionError(f"连接失败: {str(e)}")

    def disconnect_client(self, client_id: str):
        """
        断开客户端与Broker的连接并释放资源

        :param client_id: 要断开连接的客户端ID
        """
        if client_id not in self.clients:
            return

        client = self.clients[client_id]
        try:
            # 停止后台循环并断开连接
            client.loop_stop()
            client.disconnect()
            time.sleep(0.1)  # 确保断开操作完成
        finally:
            # 清理客户端相关资源
            del self.clients[client_id]
            if client_id in self.client_status:
                del self.client_status[client_id]

    def get_client_status(self, client_id: str) -> bool:
        """
        获取客户端的连接状态

        :param client_id: 客户端ID
        :return: 已连接返回True，否则返回False
        """
        return self.client_status.get(client_id, False)

    def subscribe(self, client_id: str, topic: str, qos: int = 0, timeout: int = 5) -> bool:
        """
        客户端订阅指定主题

        :param client_id: 已连接的客户端ID
        :param topic: 要订阅的主题
        :param qos: 订阅的QoS级别（0/1/2），默认值为0
        :param timeout: 订阅确认超时时间（秒），默认值为5
        :return: 订阅成功返回True，否则返回False
        :raises ConnectionError: 客户端未连接时抛出
        """
        if not self.get_client_status(client_id):
            raise ConnectionError(f"客户端 '{client_id}' 未连接，请先调用 connect_client 方法")

        client = self.clients[client_id]
        start_time = time.time()

        try:
            # 发起订阅请求
            result, mid = client.subscribe(topic, qos)
            if result != 0:
                return False  # 订阅请求发送失败

            # QoS 0无需等待确认，直接返回成功
            if qos == 0:
                return True
            # QoS 1/2等待确认（简化处理，使用超时判断）
            else:
                while time.time() - start_time < timeout:
                    time.sleep(0.1)
                return True

        except Exception as e:
            print(f"订阅失败: {str(e)}")
            return False

    def unsubscribe(self, client_id: str, topic: str, timeout: int = 5) -> bool:
        """
        客户端取消订阅指定主题（取消订阅不涉及QoS级别）

        :param client_id: 已连接的客户端ID
        :param topic: 要取消订阅的主题
        :param timeout: 取消订阅确认超时时间（秒），默认值为5
        :return: 取消订阅成功返回True，否则返回False
        :raises ConnectionError: 客户端未连接时抛出
        """
        if not self.get_client_status(client_id):
            raise ConnectionError(f"客户端 '{client_id}' 未连接，请先调用 connect_client 方法")

        client = self.clients[client_id]
        start_time = time.time()

        try:
            # 发起取消订阅请求
            result, mid = client.unsubscribe(topic)
            if result != 0:
                return False  # 取消订阅请求发送失败

            # 等待确认
            while time.time() - start_time < timeout:
                time.sleep(0.1)
            return True

        except Exception as e:
            print(f"取消订阅失败: {str(e)}")
            return False
        
    def publish_sync(self, client_id: str, topic: str, payload: str, qos: int = 0,
                     retain: bool = False, timeout: int = 10) -> bool:
        """
        同步发布消息（阻塞等待发布结果）

        :param client_id: 已连接的客户端ID
        :param topic: 发布消息的主题
        :param payload: 消息内容（字符串）
        :param qos: 发布的QoS级别（0/1/2），默认值为0
        :param retain: 是否保留消息，默认值为False
        :param timeout: 发布确认超时时间（秒），默认值为10
        :return: 发布成功返回True，否则返回False
        :raises ConnectionError: 客户端未连接时抛出
        """
        if not self.get_client_status(client_id):
            raise ConnectionError(f"客户端 '{client_id}' 未连接，请先调用 connect_client 方法")

        client = self.clients[client_id]
        try:
            # 发布消息并等待确认
            result = client.publish(topic, payload, qos, retain)
            result.wait_for_publish(timeout)
            return result.rc == 0
        except Exception as e:
            print(f"同步发布失败: {str(e)}")
            return False

    def publish_async(self, client_id: str, topic: str, payload: str, qos: int = 0,
                      retain: bool = False, timeout: int = 10) -> None:
        """
        异步发布消息（不阻塞主线程）

        :param client_id: 已连接的客户端ID
        :param topic: 发布消息的主题
        :param payload: 消息内容（字符串）
        :param qos: 发布的QoS级别（0/1/2），默认值为0
        :param retain: 是否保留消息，默认值为False
        :param timeout: 发布确认超时时间（秒），默认值为10
        """
        def _async_pub_task():
            """异步发布的子线程任务"""
            success = False
            error_msg = ""
            try:
                if not self.get_client_status(client_id):
                    error_msg = f"客户端 '{client_id}' 未连接"
                    return

                client = self.clients[client_id]
                result = client.publish(topic, payload, qos, retain)
                result.wait_for_publish(timeout)
                success = result.rc == 0
                if not success:
                    error_msg = f"发布失败（返回码: {result.rc}）"

            except Exception as e:
                error_msg = f"异步发布异常: {str(e)}"
            finally:
                # 触发回调函数（如果已注册）
                if self.async_pub_callback:
                    self.async_pub_callback(client_id, topic, success, error_msg)
                # 从线程列表中移除已完成的线程
                with self.async_pub_threads_lock:
                    if pub_thread in self.async_pub_threads:
                        self.async_pub_threads.remove(pub_thread)

        # 创建并启动异步线程
        pub_thread = threading.Thread(target=_async_pub_task, daemon=True)
        pub_thread.start()
        with self.async_pub_threads_lock:
            self.async_pub_threads.append(pub_thread)

    def set_async_pub_callback(self, callback: Callable[[str, str, bool, str], None]) -> None:
        """
        注册异步发布结果的回调函数

        :param callback: 回调函数，参数为：
                        (client_id: 客户端ID, 
                         topic: 发布的主题, 
                         success: 是否成功, 
                         error_msg: 错误信息)
        """
        self.async_pub_callback = callback

    def wait_for_all_async_pub(self, timeout: int = 30) -> None:
        """
        等待所有异步发布线程完成

        :param timeout: 等待超时时间（秒），默认值为30
        """
        start_time = time.time()
        while self.async_pub_threads and (time.time() - start_time) < timeout:
            # 清理已完成的线程
            with self.async_pub_threads_lock:
                self.async_pub_threads = [t for t in self.async_pub_threads if t.is_alive()]
            time.sleep(0.1)

    def get_received_msg(self, topic: Optional[str] = None, since: Optional[float] = None) -> Dict[str, List[Tuple[str, float]]]:
        """
        获取缓存的接收消息

        :param topic: 要获取消息的主题，None表示获取所有主题，默认值为None
        :param since: 时间戳，只返回该时间戳之后的消息，None表示返回所有消息，默认值为None
        :return: 消息字典，格式为 {topic: [(payload, timestamp), ...]}
        """
        with self.msg_cache_lock:
            result = {}
            target_topics = [topic] if topic else self.msg_cache.keys()

            for t in target_topics:
                if t not in self.msg_cache:
                    result[t] = []
                    continue

                current_time = time.time()
                filtered = []
                for payload, ts in self.msg_cache[t]:
                    # 过滤过期消息
                    if self.msg_cache_expire > 0 and (current_time - ts) > self.msg_cache_expire:
                        continue
                    # 过滤时间范围外的消息
                    if since and ts < since:
                        continue
                    filtered.append((payload, ts))

                result[t] = filtered
            return result

    def wait_for_msg(self, topic: str, timeout: int = 10) -> Optional[Tuple[str, float]]:
        """
        阻塞等待指定主题的消息

        :param topic: 要等待的主题
        :param timeout: 等待超时时间（秒），默认值为10
        :return: 收到消息返回 (payload, timestamp)，超时返回None
        """
        start_time = time.time()
        while time.time() - start_time < timeout:
            with self.msg_cache_lock:
                if topic in self.msg_cache and len(self.msg_cache[topic]) > 0:
                    return self.msg_cache[topic].pop(0)  # 返回最早收到的消息
            time.sleep(0.1)
        return None

    def clear_msg_cache(self, topic: Optional[str] = None) -> None:
        """
        清理消息缓存

        :param topic: 要清理的主题，None表示清理所有主题，默认值为None
        """
        with self.msg_cache_lock:
            if topic and topic in self.msg_cache:
                del self.msg_cache[topic]
            else:
                self.msg_cache.clear()

    def get_msg_cache_status(self) -> Dict[str, Any]:
        """
        获取当前消息缓存的状态信息

        :return: 包含缓存配置、总消息数和各主题统计的字典
        """
        with self.msg_cache_lock:
            total_count = sum(len(msgs) for msgs in self.msg_cache.values())
            topic_stats = {
                topic: {
                    "msg_count": len(msgs),
                    "oldest_msg_time": min(ts for _, ts in msgs) if msgs else None,
                    "newest_msg_time": max(ts for _, ts in msgs) if msgs else None
                }
                for topic, msgs in self.msg_cache.items()
            }

            return {
                "cache_config": {
                    "max_size": self.msg_cache_max_size,
                    "expire_seconds": self.msg_cache_expire
                },
                "total_msg_count": total_count,
                "topic_stats": topic_stats
            }

    def cleanup_all_resources(self, wait_timeout: int = 30) -> None:
        """
        清理所有资源，包括：
        1. 等待所有异步发布线程完成
        2. 断开所有客户端连接
        3. 清理所有消息缓存
        4. 重置回调函数

        :param wait_timeout: 等待异步线程完成的超时时间（秒），默认值为30
        """
        # 1. 等待所有异步发布线程完成
        self.wait_for_all_async_pub(wait_timeout)
        
        # 2. 断开所有客户端连接
        client_ids = list(self.clients.keys())  # 创建副本避免迭代中修改
        for client_id in client_ids:
            self.disconnect_client(client_id)
        
        # 3. 清理所有消息缓存
        self.clear_msg_cache()
        
        # 4. 重置回调函数
        self.async_pub_callback = None
        
        # 5. 确保线程列表为空
        with self.async_pub_threads_lock:
            self.async_pub_threads.clear()

    def create_random_msg(self):
        num = random.randint(1, 10)  
        num_str = str(num)  
        return f"{num_str}"



# 调用示例
if __name__ == "__main__":
    # 初始化工具库
    mqtt_lib = MQTTTestLib(
        broker_host="localhost",
        broker_port=1883,
        msg_cache_max_size=100,
        msg_cache_expire=30
    )

    # 注册异步发布回调
    def async_callback(client_id, topic, success, error):
        print(f"异步发布结果: 客户端{client_id} 主题{topic} {'成功' if success else '失败'} {error}")
    mqtt_lib.set_async_pub_callback(async_callback)

    # 创建并连接客户端
    pub_client_id = "publisher"
    sub_client_id = "subscriber"
    
    mqtt_lib.create_client(pub_client_id)
    mqtt_lib.create_client(sub_client_id)
    
    mqtt_lib.connect_client(pub_client_id)
    mqtt_lib.connect_client(sub_client_id)

    # 订阅主题
    test_topic = "abc/ok"
    mqtt_lib.subscribe(sub_client_id, test_topic, qos=1)
    time.sleep(1)  # 等待订阅生效

    # 发布消息
    # mqtt_lib.publish_sync(pub_client_id, test_topic, "同步消息", qos=1)
    for i in range(500):
        time.sleep(0.5)
        payload = mqtt_lib.create_random_msg()
        mqtt_lib.publish_async(pub_client_id, test_topic, payload, qos=0)

    # 等待异步发布完成并接收消息
    mqtt_lib.wait_for_all_async_pub()
    time.sleep(2)  # 等待消息接收

    # 查看接收的消息
    # print("接收的消息:", mqtt_lib.get_received_msg(topic=test_topic))

    # 清理所有资源（替代之前的分步清理）
    mqtt_lib.cleanup_all_resources()
